-module(emqx_bridge_kafka).

-include_lib("emqx/include/emqx.hrl").
-include_lib("emqx/include/logger.hrl").

-import(string, [concat/2]).
-import(lists, [nth/2]).

-export([ load/1
        , unload/0
        ]).

%% Client Lifecircle Hooks
-export([ on_client_connected/3
        , on_client_disconnected/4
        ]).

%% Message Pubsub Hooks
-export([ on_message_publish/2
        , on_message_delivered/3
        ]).



load(Env) ->
    ekaf_init([Env]),
    emqx:hook('client.connected',    {?MODULE, on_client_connected, [Env]}),
    emqx:hook('client.disconnected', {?MODULE, on_client_disconnected, [Env]}),
    emqx:hook('message.publish',     {?MODULE, on_message_publish, [Env]}),
    emqx:hook('message.delivered',   {?MODULE, on_message_delivered, [Env]}).

%%--------------------------------------------------------------------
%% Client Lifecircle Hooks
%%--------------------------------------------------------------------

on_client_connected(Client = #{
    clientid := ClientId,
    peerhost := {IPA, IPB, IPC, IPD},
    username := Username}, _Conn = #{
    connected_at := ConnectedAt
    }, _Env) ->
    IPAddress = io_lib:format("~p\.~p\.~p\.~p", [IPA, IPB, IPC, IPD]),
    {ok, KafkaTopic} = application:get_env(emqx_bridge_kafka, values),
    ProduceTopic = proplists:get_value(kafka_linkstatus_topic, KafkaTopic),
    Json = jsx:encode([
            {type,<<"connected">>},
            {clientid,ClientId},
            {clientip,unicode:characters_to_binary(IPAddress)},
            {broker_node,node()},
            {conn_servts,ConnectedAt}
    ]),
    ekaf:produce_async(ProduceTopic, Json),
    {ok, Client}.

on_client_disconnected(_Client = #{
        clientid := ClientId,
        peerhost := {IPA, IPB, IPC, IPD},
        username := Username
    }, Reason,
        _Conn = #{
        connected_at := ConnectedAt,
        disconnected_at := DisconnectedAt
    }, _Env) ->
    IPAddress = io_lib:format("~p\.~p\.~p\.~p", [IPA, IPB, IPC, IPD]),
    {ok, KafkaTopic} = application:get_env(emqx_bridge_kafka, values),
    ProduceTopic = proplists:get_value(kafka_linkstatus_topic, KafkaTopic),
    Json = jsx:encode([
            {type,<<"disconnected">>},
            {clientid,ClientId},
            {clientip,unicode:characters_to_binary(IPAddress)},
            {borker_node,node()},
            {reason,Reason},
            {conn_servts,ConnectedAt},
            {disconn_servts,DisconnectedAt}
    ]),
    ekaf:produce_async(ProduceTopic, Json),
    {ok, _Client}.

%%--------------------------------------------------------------------
%% Message PubSub Hooks
%%--------------------------------------------------------------------

on_message_publish(Message = #message{topic = <<"$SYS/", _/binary>>}, _Env) ->
    {ok, Message};

on_message_publish(Message = #message{
  id = MessageId,
  from = ClientId,
  qos = Qos,
  topic = Topic,
  payload = Payload,
  flags = #{dup := Dup, retain := Retain},
  headers = #{username := UserName},
  timestamp = Timestamp
}, _Env) ->
    {ok, KafkaTopic} = application:get_env(emqx_bridge_kafka, values),
    ProduceTopic = proplists:get_value(kafka_publish_topic, KafkaTopic),
    Json = jsx:encode([
            {type,<<"published">>},
            {topic,Topic},
            {payload,Payload},
            {clientid,ClientId},
            {msgid,hex:bin_to_hexstr(MessageId)},
            {dup,Dup},
            {qos,Qos},
            {retain, Retain},
            {broker_node,node()},
            {servts,Timestamp}
    ]),
    ekaf:produce_async(ProduceTopic, Json),
    {ok, Message}.


on_message_delivered(_ClientInfo = #{clientid := DesClientId},Message = #message{
  id = MessageId,
  from = SourceClientId,
  qos = Qos,
  topic = Topic,
  payload = Payload,
  flags = #{dup := Dup, retain := Retain},
  headers = #{username := UserName},
  timestamp = Timestamp
}, _Env) ->
    {ok, KafkaTopic} = application:get_env(emqx_bridge_kafka, values),
    ProduceTopic = proplists:get_value(kafka_delivered_topic, KafkaTopic),
    Json = jsx:encode([
            {type,<<"delivered">>},
            {topic,Topic},
            {payload,Payload},
            {sou_clientid,SourceClientId},
            {des_clientid,DesClientId},
            {msgid,hex:bin_to_hexstr(MessageId)},
            {dup,Dup},
            {qos,Qos},
            {retain, Retain},
            {broker_node,node()},
            {servts,Timestamp}
    ]),
    ekaf:produce_async(ProduceTopic, Json),
    {ok, Message}.

%% Called when the plugin application stop
unload() ->
    emqx:unhook('client.connected',    {?MODULE, on_client_connected}),
    emqx:unhook('client.disconnected', {?MODULE, on_client_disconnected}),
    emqx:unhook('message.publish',     {?MODULE, on_message_publish}),
    emqx:unhook('message.delivered',   {?MODULE, on_message_delivered}).

%% Init kafka server parameters
ekaf_init(_Env) ->
    application:load(ekaf),
    {ok, Values} = application:get_env(emqx_bridge_kafka, values),
    BootstrapBroker = proplists:get_value(bootstrap_broker, Values),
    PartitionStrategy= proplists:get_value(partition_strategy, Values),
    application:set_env(ekaf, ekaf_partition_strategy, PartitionStrategy),
    application:set_env(ekaf, ekaf_bootstrap_broker, BootstrapBroker),
    {ok, _} = application:ensure_all_started(ekaf),
    io:format("Initialized ekaf with ~p~n", [BootstrapBroker]).        



